﻿#pragma once
#include "socks_client.hpp"

extern "C"{
#include <lwip/tcp.h>
#include <lwip/netif.h>
#include <lwip/init.h>
#include <lwip/udp.h>
#include <lwip/sys.h>
#include  <lwip/timeouts.h>
}

#include "asio.hpp"
#include <iostream>
#include <set>
#include <memory>
#include <atomic>
#include <thread>
#include <chrono>
#include <timeapi.h>
#include <windivert.h>

#include "driver2socks.h"
#include "sys_arch.h"
#include "spsc_queue.h"
#include "mpmc_queue.h"
#include "netpacket_pool.h"
#include "base/xlog.h"
#include "concurrentqueue.h"
#include "app_config.h"

extern const Driver2SocksConfig* g_config;

namespace driver2socks {
	using namespace std::chrono_literals;

	class LWIPStack {
	public:
        enum EventType {

            TCP_WRITE = 1,
            TCP_CLOSE,
            IP_INPUT,
        };

        struct StackEvent {
            typedef std::shared_ptr<StackEvent> Ptr;
            void* pcb = nullptr;
            NetPacket::Ptr netpack;
            EventType type;
        };
        static LWIPStack& getInstance() {
			static LWIPStack _stack;
            return _stack;
		}

		inline static tcp_pcb* lwip_tcp_listen(tcp_pcb* pcb) {
			return tcp_listen(pcb);
		}

		inline static err_t lwip_udp_connect(struct udp_pcb *pcb, const ip_addr_t *ipaddr, u16_t port) {
			return udp_connect(pcb, ipaddr, port);
		}

		inline static void lwip_tcp_arg(tcp_pcb* pcb, void* arg) {
			return tcp_arg(pcb, arg);
		}

        inline static void lwip_tcp_receive(struct tcp_pcb* pcb, std::function<std::remove_pointer<tcp_recv_fn>::type> receive) {
            if (receive == NULL) {
                tcp_err(pcb, NULL);
            }
            else {
                tcp_err(pcb, [pcb, receive](void *arg, err_t err) {
                    receive(arg, pcb, NULL, err);
                });
            }
            return tcp_recv(pcb, receive);
        }

		inline static void lwip_tcp_recved(struct tcp_pcb* pcb, u16_t len) {
			return tcp_recved(pcb, len);
		}

		inline static void lwip_udp_timeout(struct udp_pcb* pcb, std::function<std::remove_pointer_t<udp_timeout_fn>> timeout_fn) {
			return udp_timeout(pcb, timeout_fn);
		}

		inline static void lwip_udp_create(std::function<std::remove_pointer_t<udp_crt_fn>> create_fn) {
			return udp_create(create_fn);
		}

		inline static void lwip_udp_set_timeout(udp_pcb* pcb, u32_t timeout) {
			return udp_set_timeout(pcb, timeout);
		}

		inline static void lwip_udp_recv(struct udp_pcb *pcb, std::function<std::remove_pointer<udp_recv_fn>::type> recv) {
			return udp_recv(pcb, recv, NULL);
		}

		inline static void lwip_udp_remove(struct udp_pcb* pcb) {
			return udp_remove(pcb);
		}

		inline static tcp_pcb* tcp_listen_any() {
            auto pcb = tcp_new();
			auto any = ip_addr_any;
            tcp_bind(pcb, &any, 0);
			return lwip_tcp_listen(pcb);
		}

		inline static udp_pcb* udp_listen_any() {
            auto pcb = udp_new();
			auto any = ip_addr_any;
            udp_bind(pcb, &any, 0);
			return pcb;
		}

		inline static err_t lwip_udp_send(struct udp_pcb *pcb, struct pbuf *p) {
			return udp_send(pcb, p);
		}

        // inline static u32_t lwip_tcp_sndbuf(tcp_pcb* pcb) {
        // 	return tcp_sndbuf(pcb);
        // }

        // inline static err_t lwip_tcp_output(tcp_pcb* pcb) {
        // 	return tcp_output(pcb);
        // }

        err_t lwip_tcp_close(tcp_pcb* pcb);



        void postTcpWrite(struct tcp_pcb *pcb, std::shared_ptr<NetPacket> buff);
        bool doWrite(tcp_pcb* pcb);

        inline void strand_post(std::function<void()> cb) {
			_strand->post([cb]() {
                if (cb != NULL)
                    cb();
            });
        }

        void strand_ip_input(std::shared_ptr<NetPacket> bf/*,size_t size, std::function<void(err_t)> cb*/)
		{
            StackEvent::Ptr e = std::make_shared<StackEvent>();
            e->netpack = bf;
            e->type = IP_INPUT;
            while(!queue_event_.try_push(e)) {
                queue_event_nonfull_.clear();
                queue_event_nonfull_.wait(false,std::memory_order_relaxed);
            }
            return;
            // _strand->post([=]() {
            // 	pbuf* p = pbuf_alloc(PBUF_RAW, size, PBUF_POOL);
   //              pbuf_take(p, bf->data, bf->data_len);
            // 	auto err = _loopback->input(p, _loopback);
            // 	if (cb != NULL)
            // 		cb(err);
            // 	});
		}

         void postTcpClose(tcp_pcb* pcb) {
            StackEvent::Ptr e = std::make_shared<StackEvent>();
            e->pcb = pcb;
            e->type = TCP_CLOSE;

            while(!queue_event_.try_push(e)) {
                queue_event_nonfull_.clear();
                queue_event_nonfull_.wait(false,std::memory_order_relaxed);
            }
		}

		inline void strand_udp_remove(udp_pcb* pcb) {
			_strand->post([=]() {
				LWIPStack::lwip_udp_remove(pcb);
			});
		}

		inline void strand_tcp_recved(tcp_pcb* pcb, u16_t len) {
			_strand->post([=]() {
				LWIPStack::lwip_tcp_recved(pcb, len);
			});
		}

		inline void strand_udp_send(udp_pcb* pcb, std::shared_ptr<pbuf> p, std::function<void(err_t)> cb) {
			_strand->post([=]() {
				auto err = LWIPStack::lwip_udp_send(pcb, p.get());
				if (cb != NULL)
					cb(err);
			});
		}

        inline void set_output_function(std::function<void (NetPacket*)> f) {
            ipv4outcb = f;
		}
        inline void set_outputv6_function(std::function<void (NetPacket*)> f)
		{
            ipv6outcb = f;
		}
        void setTcpSentCallback(tcp_pcb* pcb)
        {
            tcp_sent(pcb,std::bind(&LWIPStack::onTcpSent,this,std::placeholders::_1,std::placeholders::_2,std::placeholders::_3));
        }
        void setTcpRecvCallback(tcp_pcb* pcb)
        {
            tcp_recv(pcb,std::bind(&LWIPStack::onTcpRecv,this,
                                       std::placeholders::_1,std::placeholders::_2,std::placeholders::_3,std::placeholders::_4));
        }

        ~LWIPStack()
        {
            is_close_.test_and_set(std::memory_order_acquire);
            is_close_.wait(true,std::memory_order_relaxed);
		}
        inline void init(/*asio::io_context& ctx*/) {
            lwip_init();
            // _strand = new asio::io_context::strand(ctx);
            netif_default = netif_list;
            _loopback = netif_list;

            tcp_ = LWIPStack::tcp_listen_any();
            udp_ = LWIPStack::udp_listen_any();

            _loopback->output_ip6 = std::bind(&LWIPStack::ipv6OutputCallback,this,
                                              std::placeholders::_1,std::placeholders::_2,std::placeholders::_3);
            _loopback->output = std::bind(&LWIPStack::ipv4OutputCallback,this,
                                          std::placeholders::_1,std::placeholders::_2,std::placeholders::_3);
            tcp_accept(tcp_,std::bind(&LWIPStack::onTcpAccept,this,
                                       std::placeholders::_1,std::placeholders::_2,std::placeholders::_3));

            std::thread([this]() {
                while (1) {
                    if (is_close_.test()) {
                        is_close_.clear();
                        is_close_.notify_all();
                        return;
                    }
                    sys_check_timeouts();
                    bool need_sleep = true;


                    if (!tcp_pcb_wait_write_.empty()) {
                        need_sleep = false;
                        for (auto iter = tcp_pcb_wait_write_.begin(); iter!= tcp_pcb_wait_write_.end();) {
                            if (doWrite(*iter)) {
                                auto riter = iter;
                                ++iter;
                                tcp_pcb_wait_write_.erase(riter);
                            } else {
                                ++iter;
                            }
                        }
                    }

                    StackEvent::Ptr e;
                    if (queue_event_.try_pop(e) && e) {
                        need_sleep = false;
                        if (e->type == TCP_WRITE) {
                            if (!doWrite((tcp_pcb*)e->pcb)) {
                                tcp_pcb_wait_write_.insert((tcp_pcb*)e->pcb);
                            }
                        } else if (e->type == TCP_CLOSE) {
                            lwip_tcp_close((tcp_pcb*)e->pcb);
                        } else if (e->type == IP_INPUT) {
                            pbuf* p = pbuf_alloc(PBUF_RAW, e->netpack->data_len, PBUF_POOL);
                            pbuf_take(p, e->netpack->data, e->netpack->data_len);
                            auto err = _loopback->input(p, _loopback);
                            if (err != ERR_OK) {
                                SERROR("ip input error:{}",(uint32_t)err);
                            }
                        }
                        queue_event_nonfull_.test_and_set(std::memory_order_acquire);
                        queue_event_nonfull_.notify_all();
                    }
                    if (need_sleep) {
                        // timeBeginPeriod(1);
                        Sleep(50);
                        // timeEndPeriod(1);
                    }

                }
            }).detach();
        }
        /*!
        * The callback function will be passed a NULL pbuf to
        * indicate that the remote host has closed the connection. If the
        * callback function returns ERR_OK or ERR_ABRT it must have
        * freed the pbuf, otherwise it must not have freed it.
        */
        err_t onTcpRecv(void* arg, struct tcp_pcb *tpcb, struct pbuf *p, err_t err);
        //虚拟tcp接收到连接成功回调,构造一个socks5客户端然后连接到目标
        err_t onTcpAccept(void *arg, struct tcp_pcb *newpcb, err_t err);
        err_t ipv6OutputCallback(struct netif* netif, struct pbuf* p, const ip6_addr_t* ipaddr)
        {
            if (ipv6outcb) {
                NetPacket* buffer = _NetPacketPool->getPacket(p->tot_len);
                pbuf_copy_partial(p, buffer->data, p->tot_len, 0);
                ipv6outcb(buffer);
            }
            return ERR_OK;
        }

        err_t ipv4OutputCallback(struct netif *netif, struct pbuf *p, const ip4_addr_t *ipaddr)
        {
            if (ipv4outcb) {
                NetPacket* buffer = _NetPacketPool->getPacket(p->tot_len);
                pbuf_copy_partial(p, buffer->data, p->tot_len, 0);
                ipv4outcb(buffer);
            }
            return ERR_OK;
        }
        //从队列中pop已发送数据
        err_t onTcpSent(void *arg, struct tcp_pcb *tpcb, u16_t len);
	private:
        LWIPStack() : _strand(NULL), _loopback(NULL), queue_event_{4096} {}
        //ipv6回调
        std::function<void (NetPacket*)> ipv6outcb;
        //ipv4
        std::function<void (NetPacket*)> ipv4outcb;
        tcp_pcb* tcp_;
        udp_pcb* udp_;
        std::atomic_uint32_t total_{0};
        std::set<tcp_pcb*> tcp_pcb_wait_write_;

		asio::io_context::strand* _strand;
		netif* _loopback;
		std::thread thread_stack_;
        MPMCQueue<StackEvent::Ptr> queue_event_;
        std::atomic_flag queue_event_nonfull_{};
        std::atomic_flag is_close_{};
	};
}
